%% Load the data - Select the File in 'ViscotaxisExperiments'
close all; clear all; clc;
clearvars;
total_directory=uigetdir();
total_directory = dir(total_directory);
dirFlags = [total_directory.isdir];
total_directory = total_directory(~ismember({total_directory.name},{'.','..'}));
cd([total_directory(1).folder])
%% Combine Tracks
tracks_ExpComb = []; c = 0;h = waitbar(0,'Calculating ... ');
for i = 1:length(total_directory)
    clear tracks
    load([total_directory(i).folder filesep total_directory(i).name]);
    tracks_ExpComb{i} = tracks;
end
waitbar(c+1 / 5,h);
% Filter Tracks
TracksSel = [];YwallMax_s = [];YwallMin_s = [];myparameters = [];
myparameters.lengthdev = 4;

for i = 1:length(tracks_ExpComb)
    clear InputTracks tracks_c tracks_filt
    InputTracks = tracks_ExpComb{1,i};
    [tracks_c] = Camera2MicronSecond(InputTracks,mag,fps,Camera_um_pix);
    [tracks_filt, trackstats] = FilterTracks(tracks_c, myparameters);

    YwallMin_s(1,i) = min(vertcat(tracks_filt(:).y_um));
    YwallMax_s(1,i) = max(vertcat(tracks_filt(:).y_um));
    TracksSel{i} = tracks_filt;
end
waitbar(c+1 / 5,h);
% Calculate Spatial PDF and Swimming Speed
Nbins = 11;
YwallMin = min(YwallMin_s); YwallMax = max(YwallMax_s);

results_Exp = [];
for i = 1:length(TracksSel)
    clear results InputTracks
    InputTracks = TracksSel{1,i};
    [results, ListYs, BinSize] = SpeedandPDF_Final(InputTracks, fps, mag, Nbins,YwallMin,YwallMax);
    results_Exp{i} = results
end
waitbar(c+1 / 5,h);

if strcmp(PerPEO,'0.5');RevExp = 2; else; RevExp = 0; end
[results,FinalResults] = ReplicateExpCombine(results_Exp,RevExp);
waitbar(c+1 / 5,h); close(h)
'Done'
%% Plot
close all;
Marker = '-';
FillColor = [0.25 0.25 0.25]; LineColor = [0 0 0];
LineWidth = 1;
alpha_errorbar = 0.50;
alw = 0.5; % Border line width
DensityorSpeed = 'Density'; %write 'Density' or 'Speed'
DensityorSpeed = 'Speed'

[fillErrorBar,x,y,error] = SpatialSpeedDensityPlot(DensityorSpeed,FinalResults,LineWidth,alpha_errorbar,PerPEO)

xlabel('Position in y [micron]','FontSize',14);
ylabel(DensityorSpeed,'FontSize',14);

% Info for combine PDF curve
mypapersize = [.9 .9]; %[width height]

if strcmp(DensityorSpeed, 'Density')
    SingleCurvePDF = gcf;
    set(gca,'LineWidth',alw);
    set(gca,'XColor',[0 0 0]);set(gca,'YColor',[0 0 0]);
    set(gcf, 'PaperUnits', 'inches', 'papersize', mypapersize, 'PaperPosition', [0 0 mypapersize]);
    set(SingleCurvePDF,'Units','inches','Position',[0 0 mypapersize]);

    xticks([0,250,500,750,1000]); xticklabels({'','','','',''});
    yticks([0,0.5E-3,1.0E-3,1.5E-3]); yticklabels({'','','',''});
    ylabel('');xlabel('');title('')
    ax = gca;outerpos = ax.OuterPosition;ti = ax.TightInset;
    left = outerpos(1);bottom = outerpos(2);ax_width = outerpos(3);ax_height = outerpos(4);
    ax.Position = [left+0.011 bottom+0.01 ax_width-0.025 ax_height-0.022];

   figdir = '.\ExtendedDataFig6';
   print(SingleCurvePDF, '-painters','-dpdf', '-r1200', [figdir,'\MeanPDF',PerPEO,'.pdf']); %r1200
end
% Info for Combine Speed Curve
clc;
BorderColor = [0 0 0];
if strcmp(DensityorSpeed, 'Speed')
    SingleCurveSpeed = gcf;
    set(gca,'LineWidth',alw);set(gca,'XColor',BorderColor);set(gca,'YColor',BorderColor);
    set(gcf, 'PaperUnits', 'inches', 'papersize', mypapersize, 'PaperPosition', [0 0 mypapersize]);
    set(SingleCurveSpeed,'Units','inches','Position',[0 0 mypapersize]);
    xticks([0,250,500,750,1000]); xticklabels({'','','','',''});
    yticks([0,25,50,75,100]); yticklabels({'','','',''});
    ylabel('');xlabel('');title('');
    ax = gca;outerpos = ax.OuterPosition;ti = ax.TightInset;left = outerpos(1);bottom = outerpos(2);ax_width = outerpos(3);ax_height = outerpos(4);
    ax.Position = [left+0.011 bottom+0.01 ax_width-0.025 ax_height-0.022];
    
   figdir = '.\ExtendedDataFig6';  
   print(SingleCurveSpeed, '-painters','-dpdf', '-r1200', [figdir,'\MeanSpeed_',PerPEO,'.pdf']);      
end

% Save Data in Excel
dir_Excel = '.\Source data\ExtendedData\'
if strcmp(DensityorSpeed, 'Density')
    Tab = 1;
    if strcmp(PerPEO,'0.05'); column1 = 'E4'; column2 = 'F4'; column3 = 'G4'; 
    elseif strcmp(PerPEO,'0.00'); column1 = 'A4'; column2 = 'B4'; column3 = 'C4'; 
    else; column1 = []; column2 = []; column3 = [];
    end
    writematrix(x',[dir_Excel,'ExtDataFigure6.xlsx'],'Sheet',Tab,'Range',column1)
    writematrix(y',[dir_Excel,'ExtDataFigure6.xlsx'],'Sheet',Tab,'Range',column2)
    writematrix(error',[dir_Excel,'ExtDataFigure6.xlsx'],'Sheet',Tab,'Range',column3)
end

if strcmp(DensityorSpeed, 'Speed') 
    Tab = 2;
    if strcmp(PerPEO,'0.05'); column1 = 'E4'; column2 = 'F4'; column3 = 'G4'; 
    elseif strcmp(PerPEO,'0.00'); column1 = 'A4'; column2 = 'B4'; column3 = 'C4'; 
    else; column1 = []; column2 = []; column3 = [];
    end
    writematrix(x',[dir_Excel,'ExtDataFigure6.xlsx'],'Sheet',Tab,'Range',column1)
    writematrix(y',[dir_Excel,'ExtDataFigure6.xlsx'],'Sheet',Tab,'Range',column2)
    writematrix(error',[dir_Excel,'ExtDataFigure6.xlsx'],'Sheet',Tab,'Range',column3)
end
%% FUNCTION: Convert pix/frame to um/s
function [c_tracks] = Camera2MicronSecond(tracks,mag,fps,Camera_micron_pix)
% Camera2MicronSecond: Converts camera units (frames and pix) to physical
% units (microns and seconds). Cameras have different pixel dimensions. 
% For example, the Zyla has 6.5 microns per pixel
    c_tracks = tracks;
    for i=1:length(c_tracks)
        c_tracks(i).t_sec = c_tracks(i).t/fps;
        c_tracks(i).x_um = c_tracks(i).x.*Camera_micron_pix/mag; c_tracks(i).y_um = c_tracks(i).y.*Camera_micron_pix/mag;
        c_tracks(i).u_um = c_tracks(i).u.*Camera_micron_pix.*fps./mag; c_tracks(i).v_um = c_tracks(i).v.*Camera_micron_pix.*fps./mag;    
        % Calculate the Speed of the Trajectory (magnitude of u-velocity and v-velocity)    
        s = sqrt(c_tracks(i).v.^2+c_tracks(i).u.^2); c_tracks(i).Speed = s;
        s_um = sqrt(c_tracks(i).v_um.^2+c_tracks(i).u_um.^2);  c_tracks(i).Speed_um = s_um;   
    end
end
%% FUNCTION: Filter Tracks
function [f_tracks, trackstats] = FilterTracks(tracks, myparameters)

if ~isfield(myparameters,'minmaxlengthtime')||isempty(myparameters.minmaxlengthtime); myoptions.minmaxlengthtime=0; else; minlengthtime = myparameters.minmaxlengthtime(1);maxlengthtime=myparameters.minmaxlengthtime(2); myoptions.minmaxlengthtime=1; end
if ~isfield(myparameters,'minmedianspeed')||isempty(myparameters.minmedianspeed);  myoptions.minmedianspeed=0; else; minmedianspeed = myparameters.minmedianspeed; myoptions.minmedianspeed=1; end
if ~isfield(myparameters,'minRG') || isempty(myparameters.minRG); myoptions.minRG = 0; else; minRG = myparameters.minRG; myoptions.minRG = 1; end
if ~isfield(myparameters,'lengthdev')||isempty(myparameters.lengthdev); myoptions.lengthdev=0; else; lengthdev=myparameters.lengthdev; myoptions.lengthdev=1; end
if ~isfield(myparameters,'minmaxcontourlength')||isempty(myparameters.minmaxcontourlength);myoptions.minmaxcontourlength=0; else; mincontourlength=myparameters.minmaxcontourlength(1); maxcontourlength=myparameters.minmaxcontourlength(2); myoptions.minmaxcontourlength=1; end

f_tracks=tracks;

if myoptions.minmaxlengthtime==1
    ind_minmaxlengthtime=zeros(1,length(f_tracks));tracklength=zeros(1,length(f_tracks));
    for k=1:length(f_tracks)
        tracklength(k) = length(f_tracks(k).x);
        if tracklength(k) < maxlengthtime && tracklength(k) > minlengthtime
            ind_minmaxlengthtime(1,k)=1;
        end
    end
    trackstats.tracklengthtime=tracklength;
else
    ind_minmaxlengthtime=ones(1,length(f_tracks));
end

if myoptions.minmedianspeed==1
    med=zeros(1,length(f_tracks));ind_minmedianspeed=zeros(1,length(f_tracks));
    for k=1:length(f_tracks)
        u_r = sqrt(tracks(k).u.^2 + tracks(k).v.^2);
        med(k) = median(u_r);
        if med(k) > minmedianspeed
            ind_minmedianspeed(1,k)=1;
        end
    end
    trackstats.medianspeed=med;
else
    ind_minmedianspeed=ones(1,length(f_tracks));
end

if myoptions.minRG==1
    ind_minRG=zeros(1,length(f_tracks));RG=zeros(1,length(f_tracks));
    for k=1:length(f_tracks)
        RG(1,k)= sqrt((max(f_tracks(k).x)-min(f_tracks(k).x))^2+(max(f_tracks(k).y)-min(f_tracks(k).y))^2);
        if RG(1,k)>minRG
            ind_minRG(1,k)=1;
        end
    end
    trackstats.RG=RG;
else
    ind_minRG=ones(1,length(f_tracks));
    
end

if myoptions.lengthdev==1
    ind_lengthdev=zeros(1,length(f_tracks));RG=zeros(1,length(f_tracks));contourlength=zeros(1,length(f_tracks));dev=zeros(1,length(f_tracks));
    s_tracks=EqualSpace(f_tracks,1);
    for k=1:length(f_tracks)
        RG(1,k)= sqrt((max(f_tracks(k).x)-min(f_tracks(k).x))^2+(max(f_tracks(k).y)-min(f_tracks(k).y))^2);
        contourlength(1,k)=length(s_tracks(k).x);
        dev(1,k)=contourlength(1,k)/RG(1,k);
        if dev(1,k)<lengthdev
            ind_lengthdev(1,k)=1;
        end
    end
    trackstats.RG=RG;
    trackstats.contourlength=contourlength;
    trackstats.dev=dev;
else
    ind_lengthdev=ones(1,length(f_tracks));
    
end

if myoptions.minmaxcontourlength==1
    ind_minmaxcontourlength=zeros(1,length(f_tracks));
    if myoptions.lengthdev==0
        s_tracks=EqualSpace(f_tracks,1);contourlength=zeros(1,length(f_tracks));
        for k=1:length(s_tracks)
            contourlength(1,k)=length(s_tracks(k).x);
            if contourlength(1,k)> mincontourlength && contourlength(1,k)<maxcontourlength
                ind_minmaxcontourlength(1,k)=1;
            end
        end
        trackstats.contourlength=contourlength;
    else
        for k=1:length(contourlength)
            if contourlength(1,k)> mincontourlength && contourlength(1,k)< maxcontourlength
                ind_minmaxcontourlength(1,k)=1;
            end
        end
    end 
else
    ind_minmaxcontourlength=ones(1,length(f_tracks));
end

ind_total=ind_minmaxlengthtime & ind_minmedianspeed & ind_minRG & ind_lengthdev & ind_minmaxcontourlength;

f_tracks=f_tracks(ind_total);
end

function tracks_lambda=EqualSpace(tracks,spacing)
tracks_lambda=struct('x',[],'y',[],'s',[],'u',[],'v',[],'t',[]);
for i=1:length(tracks)
    s=zeros(1,length(tracks(i).x));
    for ii=2:length(tracks(i).x)
        s(ii)=s(ii-1)+sqrt((tracks(i).x(ii)-tracks(i).x(ii-1))^2+(tracks(i).y(ii)-tracks(i).y(ii-1))^2);
    end
    sq=0:spacing:s(end);
    x=interp1(s,tracks(i).x,sq);y=interp1(s,tracks(i).y,sq);u=interp1(s,tracks(i).u,sq,'pchip');v=interp1(s,tracks(i).v,sq,'pchip');t=interp1(s,tracks(i).t,sq);
    tracks_lambda(i,1).x(:,1)=x;tracks_lambda(i,1).y(:,1)=y;tracks_lambda(i,1).u(:,1)=u;tracks_lambda(i,1).v(:,1)=v;tracks_lambda(i,1).s(:,1)=s;tracks_lambda(i,1).t(:,1)=t;
end
end
%% FUNCTION: Bin the tracks to get Spatial Density and Swimming Speed
function [results, ListYs, BinSize] = SpeedandPDF_Final(tracks_in, fps, mag, Nbins,YwallMin,YwallMax)
%SpeedandPDF_Final: Using the tracks from the Guasto Lab tracking
%code (tracks) we can bin the y-position with corresponding speeds to
%create a PDF and speed distribution of the tracks in the channel.

    tracks = tracks_in;
    Ys= cat(1,tracks(:).y_um);  Ys_um = Ys;   
    Speed = cat(1,tracks(:).Speed_um);
    Urs = cat(1,tracks(:).Speed_um); U_um = Urs;
    
    %Determine the minumum and maximum y values
    Ymin = YwallMin; Ymax = YwallMax;

    ListYs = linspace(Ymin,Ymax,Nbins+1);   %Create the y-intervals to categorize the speeds 
    Counts  = zeros(1,Nbins);               %create the matrix to store the number of data in the section
    Speeds  = Counts;                       %Initialize the variable speeds and counts
    Speeds_um  = Counts;
    SpeedsSquare = Counts;                  %Initialize the variable SpeedsSquare and Counts
    SpeedsSquare_um = Counts;
    BinSize = ListYs(2)-ListYs(1);          %Determine the size of each bin

    for jj = 1:length(Ys)                                              % loop to run through all the y values from Ys
        Yjj = Ys(jj);                                                  % Assigns the variable Yjj to the actual y value
        Threshold = (Yjj-ListYs(1))/BinSize;                           % Determine which bin the y position falls into
        RowY = ceil(Threshold);                                        % Rounds the Threshold up to assign it to the bin number
        if (0<RowY)&&(RowY<=Nbins)                                     % Loop to check if the Ys value falls into selected bin
        Counts(RowY) = Counts(RowY)+1;                                 % Counter to track the number of data points in each bin
        Speeds(RowY) = Speeds(RowY)+Urs(jj);                           % Counter for the average calculation below (sum of the speeds for each bin)
        SpeedsSquare(RowY) = SpeedsSquare(RowY) + Urs(jj)^2;           % Counter for the Variance equation
        
        Speeds_um(RowY) = Speeds_um(RowY)+U_um(jj);                    % um/sec: Counter for the average calculation below (sum of the speeds for each bin)
        SpeedsSquare_um(RowY) = SpeedsSquare_um(RowY) + U_um(jj)^2;    % um/sec: Counter for the Variance equation
        
        end
    end

    AverageSpeeds = Speeds./Counts;     % Calculate the average speed for each bin
    AverageSpeeds_um = Speeds_um./Counts;
    
    VarianceSpeeds = SpeedsSquare./Counts-AverageSpeeds.^2;     % Calculate the variance
    VarianceSpeeds_um = SpeedsSquare_um./Counts-AverageSpeeds_um.^2;
    d = 1;
    results(d).Speeds = AverageSpeeds; results(d).Errors = VarianceSpeeds;
    results(d).sqrVariance = sqrt(VarianceSpeeds); results(d).Ybins = ListYs(1:end-1)+BinSize/2;
    results(d).InverseMeanSpeed = 1./AverageSpeeds; results(d).ErrorBar = sqrt(VarianceSpeeds)/Counts;
    results(d).Counts = Counts;
    
    [pdf,ybins] = hist(Ys,Nbins); trapz(ybins,pdf); pdf = pdf/trapz(ybins,pdf);
    results(d).PDF = pdf; results(d).YbinsPDF = ybins;
    
    [pdf_speed,ybins_s] = hist(Speed,Nbins); trapz(ybins_s,pdf_speed); pdf_speed = pdf_speed/trapz(ybins_s,pdf_speed);
    results(d).PDFSpeed = pdf_speed; results(d).YbinsSpeed = ybins_s;
    
    results(d).Speeds_um = AverageSpeeds_um; results(d).Errors_um = VarianceSpeeds_um; results(d).sqrVariance_um = sqrt(VarianceSpeeds_um);
    results(d).InverseMeanSpeed_um = 1./AverageSpeeds_um; results(d).ErrorBar_um = sqrt(VarianceSpeeds_um)/Counts;
    
    [pdf_um,ybins_um] = hist(Ys_um,Nbins); trapz(ybins_um,pdf_um);
    pdf_um = pdf_um/trapz(ybins_um,pdf_um);
    results(d).PDF_um = pdf_um; results(d).YbinsPDF_um = ybins_um;
end
%% FUNCTION: Combine replicate experiments
function [results,FinalResults] = ReplicateExpCombine(results_Exp_in,RevExp)
% ReplicateExpCombine: Combine the replicate experiments for each viscosity profile and
% calculate the mean spatial swimming speed and PDF. Return the values
% for making figures. For 5 cP viscotaxis experiment, Exp2, flip the data
% due to the gradient being the other way.

results_Exp = results_Exp_in;
results = [];FinalResults = [];
for i = 1:length(results_Exp)
    results(i).Pos = results_Exp{1,i}.Ybins;
    if i == RevExp
        results(i).Density = flip(results_Exp{1,i}.PDF_um); results(i).Speed = flip(results_Exp{1,i}.Speeds_um);
    else
        results(i).Density = results_Exp{1,i}.PDF_um; results(i).Speed = results_Exp{1,i}.Speeds_um;
    end
end
FinalResults.Pos = mean(vertcat(results.Pos));
Store_MeanExpPDF = vertcat(results.Density); FinalResults.Density = mean(Store_MeanExpPDF);
FinalResults.Density_STD = nanstd(Store_MeanExpPDF); FinalResults.Density_STDError = nanstd(Store_MeanExpPDF)./sqrt(length(results_Exp));
Store_MeanExpSpeed = vertcat(results.Speed); FinalResults.Speed = mean(Store_MeanExpSpeed);
FinalResults.Speed_STD = nanstd(Store_MeanExpSpeed); FinalResults.Speed_STDError = nanstd(Store_MeanExpSpeed)./sqrt(length(results_Exp));
end
%% FUNCTION: Plot the figure
function [fillErrorBar,x,y,error] = SpatialSpeedDensityPlot(DensityorSpeed,FinalResults,LineWidth,alpha_errorbar,PerPEO)

PerPEOColor = [1 1 1;               % Control 
                0.635 0.078 0.184;  % 0.05 % PEO
                0.588 0.286 0.612;  % 0.1% PEO
                0.929 0.694 0.125;  % 0.25% PEO
                0.466 0.674 0.188;  % 0.5% PEO
                0.85 0.325 0.0980]; % 0.82% PEO

if strcmp(PerPEO,'0.1')
    LineColor = PerPEOColor(3,:);
elseif strcmp(PerPEO,'0.25')
    LineColor = PerPEOColor(4,:);
elseif strcmp(PerPEO,'0.5')
    LineColor = PerPEOColor(5,:);
elseif strcmp(PerPEO,'0.82')
    LineColor = PerPEOColor(6,:);
elseif strcmp(PerPEO,'0.05')
    LineColor = PerPEOColor(2,:);
else
    PerPEOColor = [1 1 1]; LineColor = [0 0 0];
end

x = FinalResults.Pos; x = x - (median(x)-500);
x_sim = []; y_sim = []; FillColor = LineColor;
if strcmp(DensityorSpeed, 'Density')
    y = FinalResults.Density; error = FinalResults.Density_STDError;
    y_max = 0.002; ticks_y = [0,0.5E-3,1.0E-3,1.5E-3];
elseif strcmp(DensityorSpeed, 'Speed')
    y = FinalResults.Speed; error = FinalResults.Speed_STDError;
    y_max = 112; ticks_y = [0,25,50,75,100];
else
    'Change DensityorSpeed Variable'
    fillErrorBar = [];
    return
end
    fillErrorBar = figure('color',[1 1 1]);
        plot(x,y,'-','Color',LineColor,'LineWidth',LineWidth)
        hold on;
        patch = fill([x,fliplr(x)],[y + error,fliplr(y - error)],FillColor);
        set(patch, 'edgecolor', 'none'); set(patch, 'FaceAlpha', alpha_errorbar);      
        hold off;  xlim([0 1000]); ylim([0 y_max]); box on;
        set(gca, 'Layer', 'top'); xticks([0,250,500,750,1000]); yticks(ticks_y);
end